package com.algorithm.example.tree;


/**
 * <h3>AVL 树</h3>
 * <ul>
 *     <li>二叉搜索树在插入和删除时，节点可能失衡</li>
 *     <li>如果在插入和删除时通过旋转, 始终让二叉搜索树保持平衡, 称为自平衡的二叉搜索树</li>
 *     <li>AVL 是自平衡二叉搜索树的实现之一</li>
 * </ul>
 *
 * @author lzt
 * @version 1.0
 * @since 2024/2/14 23:41
 */
public class AVLTree {

    AVLNode root;

    static class AVLNode {
        int height = 1;
        int key;
        Object value;
        AVLNode right;
        AVLNode left;

        public AVLNode(int key, Object value) {
            this.key = key;
            this.value = value;
        }

        public AVLNode(int key) {
            this.key = key;
        }

        public AVLNode(int key, Object value, AVLNode right, AVLNode left) {
            this.key = key;
            this.value = value;
            this.right = right;
            this.left = left;
        }


        @Override
        public String toString() {
            return "\nAVLNode{" +
                    "height=" + height +
                    ", key=" + key +
                    ", value=" + value +
                    ", right=" + right +
                    ", left=" + left +
                    '}';
        }
    }

    /**
     * 求节点的高度
     */
    private int height(AVLNode node) {
        if (node == null) {
            return 0;
        }
        return node.height;
    }

    /**
     * 更新节点高度
     */
    private void updateHeight(AVLNode node) {
        while (node != null) {
            int leftHeight = height(node.left);  // 左子树高度
            int rightHeight = height(node.right);  // 右子树高度

            node.height = Math.max(leftHeight, rightHeight) + 1;  // 更新节点高度

            node = findParent(node);  // 找到当前节点的父节点
        }
    }

    // 辅助方法，找到指定节点的父节点
    private AVLNode findParent(AVLNode node) {
        if (node == root) {
            return null;  // 如果当前节点是根节点，返回null
        }

        AVLNode parent = null;
        AVLNode current = root;
        while (current != null) {
            if (current.left == node || current.right == node) {
                parent = current;
                break;
            } else if (node.key < current.key) {
                current = current.left;
            } else {
                current = current.right;
            }
        }
        return parent;
    }

    //平衡因子（balance factor）=左子树高度-右子树高度1-1θ
    private int bf(AVLNode node) {
        if (node == null) {
            return 0;
        }
        return height(node.left) - height(node.right);
    }


    // 参数：要旋转的节点，返回值：新的根节点
    private AVLNode rightRotate(AVLNode node) {
        if (node == null) {
            return null;
        }
        AVLNode left = node.left;
        node.left = left.right;
        left.right = node;
        updateHeight(left);
        updateHeight(node);
        return left;
    }

    // 参数：要旋转的节点，返回值：新的根节点
    private AVLNode leftRotate(AVLNode node) {
        if (node == null) {
            return null;
        }
        AVLNode right = node.right;
        node.right = right.left;
        right.left = node;
        updateHeight(node);
        updateHeight(right);
        return right;
    }

    //先左旋左子树，再右旋根节点
    private AVLNode leftRightRotate(AVLNode node) {
        if (node == null) {
            return null;
        }
        node.left = leftRotate(node.left);
        node = rightRotate(node);
        return node;
    }

    //先右旋右子树，再左旋根节点
    private AVLNode rightLeftRotate(AVLNode node) {
        if (node == null) {
            return null;
        }
        node.right = leftRotate(node.right);
        node = leftRotate(node);
        return node;
    }

    // 检查节点是否失衡，重新平衡代码
    private AVLNode balance(AVLNode node) {
        if (node == null) {
            return null;
        }
        int bf = bf(node);
        int bf1;
        if (bf > 1) {
            bf1 = bf(node.left);
            if (bf1 > 1) {
                node = rightRotate(node); // LL
            } else {
                node = leftRightRotate(node); // LR
            }
        } else if (bf < -1) {
            bf1 = bf(node.right);
            if (bf1 > 1) {
                node = rightLeftRotate(node); // RL
            } else {
                node = leftRotate(node); // RR
            }
        }
        return node;
    }

    /**
     * 添加
     *
     * @param value - 待添加值
     */
    public void put(int key, Object value) {
        AVLNode p = root;
        AVLNode p1 = null;
        if (root == null) {
            root = new AVLNode(key, value);
            return;
        }
        while (p != null) {
            p1 = p;
            if (key > p.key) {
                p = p.right;
            } else if (key < p.key) {
                p = p.left;
            } else {
                p.value = value;
                return;
            }
        }
        if (p1.key < key) {
            p1.right = new AVLNode(key, value);
            updateHeight(p1.right);
        } else {
            p1.left = new AVLNode(key, value);
            updateHeight(p1.left);
        }
        root = balance(root);
    }

    /**
     * 删除
     *
     * @param key - 被删键
     */
    public Object remove(int key) {
        AVLNode p = root;
        AVLNode parent = null;
        while (p != null) {
            if (p.key > key) {
                parent = p;
                p = p.left;
            } else if (p.key < key) {
                parent = p;
                p = p.right;
            } else {
                break;
            }
        }

        if (p == null) {
            return null;
        }
        // 1.左孩子
        if (p.left != null && p.right == null) {
            shift(parent, p, p.left);
            p.left = p.right = null;
        }

        // 2.右孩子
        if (p.left == null && p.right != null) {
            shift(parent, p, p.right);
            p.left = p.right = null;
        }

        // 3.两个孩子
        if (p.right != null) {
            // 被删除节点找后续
            AVLNode s = p.right;
            AVLNode parent1 = p; //后续父亲
            while (s.left != null) {
                parent1 = s;
                s = s.left;
            }

            // 删除和后继不相邻，处理后继的后事
            if (parent1 != p) {
                shift(parent1, s, s.right); // 不可有左孩子
                s.right = p.right;
            }
            // 后继取代被删除节点
            shift(parent, p, s);
            s.left = p.left;
        }

        updateHeight(p);
        return balance(p);
    }

    /**
     * 托孤方法
     *
     * @param parent  被删除节点的父亲
     * @param deleted 被删除节点
     * @param child   被顶上去的节点
     */
    private void shift(AVLNode parent, AVLNode deleted, AVLNode child) {
        if (parent == null) {
            root = child;
            return;
        }
        if (deleted == parent.left) {
            parent.left = child;
        } else {
            parent.right = child;
        }
    }


    @Override
    public String toString() {
        return "AVLTree{" +
                "root=" + root + "平缓因子" + bf(root) +
                '}';
    }
}
